package com.thermal.thermal_inspection.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executors;

@RestController
@RequestMapping("/qwen")
public class QwenStreamController {

    private final CloseableHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private static String apiKey = "sk-35f014ffe86a406d826ce55f0096d6bf";

    private static final String API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation";

    public QwenStreamController(CloseableHttpClient httpClient, ObjectMapper objectMapper) {
        this.httpClient = httpClient;
        this.objectMapper = objectMapper;
    }

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamChat(@RequestParam String question) {
        System.out.println("进入流式传输接口！");
        SseEmitter emitter = new SseEmitter(60_000L); // 60秒超时

        Executors.newSingleThreadExecutor().execute(() -> {
            try {
                HttpPost httpPost = new HttpPost(API_URL);

                // 设置请求头
                httpPost.setHeader("Content-Type", "application/json");
                httpPost.setHeader("Authorization", "Bearer " + apiKey);
                httpPost.setHeader("X-DashScope-SSE", "enable");

                // 构建请求体
                Map<String, Object> requestBody = new HashMap<>();
                requestBody.put("model", "qwen-plus");

                Map<String, String> message = new HashMap<>();
                message.put("role", "user");
                message.put("content", question);

                requestBody.put("input", Map.of("messages", List.of(message)));
                requestBody.put("parameters", Map.of("stream", true, "incremental_output", true));

                httpPost.setEntity(new StringEntity(objectMapper.writeValueAsString(requestBody)));

                // 执行请求
                CloseableHttpResponse response = httpClient.execute(httpPost);

                try (InputStream inputStream = response.getEntity().getContent();
                     BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {

                    String line;
                    while ((line = reader.readLine()) != null) {
                        if (line.startsWith("data: ")) {
                            String jsonStr = line.substring(6).trim();
                            if (jsonStr.equals("[DONE]")) {
                                emitter.complete();
                                break;
                            }

                            JsonNode jsonNode = objectMapper.readTree(jsonStr);
                            JsonNode outputNode = jsonNode.path("output");
                            if (!outputNode.isMissingNode()) {
                                String text = outputNode.path("text").asText();
                                if (text != null && !text.isEmpty()) {
                                    emitter.send(SseEmitter.event().data(text));
                                }
                            }
                        }
                    }
                }
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }
}